# -*- coding: utf-8 -*-
"""
Created on Wed Apr 14 11:05:17 2021

@author: Lenovo-pc
"""

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import tensorflow as tf
import numpy as np

class rnn_1d:
    def __init__(self):
        pass
    
    def train(self,train_data,train_labels):
        assert (len(train_data)==len(train_labels)) and\
            len(train_data)>0 and len(train_labels)>0,\
        'RNN_1D ERROR:train input'

        row_size,col_size = train_data.shape
        input_shape =(row_size,int(col_size/3),3) #样本数\编码数（通道数）\特征数
        train_data = train_data.reshape(input_shape).astype(np.float32)
        
        self.model = tf.keras.models.Sequential([
            tf.keras.layers.LSTM(units=50),
            tf.keras.layers.Dense(1)     
            ])
        
        optimizer = tf.keras.optimizers.RMSprop(0.001)
        self.model.compile(optimizer = optimizer, 
                           loss = 'mse',
                           metrics = ['mae','mse']
                           )
        self.model.fit(train_data,train_labels,epochs=5)
        
        
    
    def estimate_cor(self,test_labels,predict_test):
        cor = np.corrcoef(test_labels,predict_test)
        return cor
    
    def save_model(self,model_save_path):
        self.model.save(model_save_path)
    
    def load_model(self,model_save_path):
        self.model = tf.keras.models.load_model(model_save_path)
        
        
    def predict_data(self,input_data):
        row_size,col_size = input_data.shape
        input_shape =(row_size,int(col_size/3),3) #样本数\编码数（通道数）\特征数
        input_data = input_data.reshape(input_shape).astype(np.float32)
        pred_data = self.model.predict(input_data)
        pred_data = pred_data.reshape(len(input_data))
        return pred_data
    
if __name__ =="__main__":
    pass